import os
import sys
import cv2
import numpy as np
import matplotlib
matplotlib.use('Agg')  # 使用非交互式后端
import matplotlib.pyplot as plt
from pathlib import Path
import time
import subprocess

# 添加远程调试支持
DEBUG_MODE = os.environ.get('DEBUG_MODE', '0') == '1'

class PositionVisualizer:
    """
    用于可视化标准位置和测试位置的矩形
    
    使用YoloSegPose类的set_standard_position方法获取标准位置向量，
    使用test_one_image方法获取测试位置向量，
    返回向量结构为 - center: 转换后的中心点坐标
            - bottom_vec: 转换后的x方向向量
            - left_vec: 转换后的y方向向量
            - z_vec: 转换后的z方向向量
            - bottom_angle: bottom向量与x轴的夹角(度, 逆时针为正)
            - left_angle: left向量与y轴的夹角(度, 逆时针为正)
    以这些向量为左下角，绘制固定大小的矩形（长2.5米，宽0.8米）。
    """

    def __init__(self, seg_model_path, pose_model_path, standard_image_path, device=0, rect_size=(0.8, 2.5), scale_factor=100):
        """
        初始化PositionVisualizer类
        
        Args:
            seg_model_path: 分割模型路径
            pose_model_path: 姿态模型路径
            device: 设备ID
            rect_size: 矩形大小（宽, 长），单位为米，默认为(0.8, 2.5)
            scale_factor: 米到像素的转换因子，默认为100（1米=100像素）
        """
        # 导入YoloSegPose类（假设已实现）
        from yoloSeg import YoloSegPose
        self.yolo_seg_pose = YoloSegPose(seg_model_path, pose_model_path, device)
        self.standard_vector = self.yolo_seg_pose.set_standard_position(standard_image_path)
        self.rect_size = rect_size  # (宽度0.8米, 长度2.5米)
        self.scale_factor = scale_factor  # 米到像素的转换因子
        self.output_dir = "results_position_visualization"
        
        # 创建输出目录
        os.makedirs(self.output_dir, exist_ok=True)


    def calculate_position(self, image_path):
        """
        测试位置
        
        Args:
            image_path: 测试位置图像路径
            
        Returns:
            测试位置向量，包含以下结构：
            - center: 转换后的中心点坐标 (米)
            - bottom_vec: 转换后的x方向向量 (米)
            - left_vec: 转换后的y方向向量 (米)
            - z_vec: 转换后的z方向向量 (米)
            - bottom_angle: bottom向量与x轴的夹角(度, 逆时针为正)
            - left_angle: left向量与y轴的夹角(度, 逆时针为正)
        """
        test_vector = self.yolo_seg_pose.test_one_image(image_path)
        if test_vector:
            return test_vector
        return None
    
    def draw_rectangle(self, vector, ax, label_prefix="标准", color=(0, 1, 0), thickness=2):
        """
        使用Matplotlib在指定轴上绘制矩形、向量和角点，矩形固定为2.5米x0.8米
        
        Args:
            vector: 位置向量，包含以下结构：
                - center: 转换后的中心点坐标（矩形左下角点，米）
                - bottom_vec: 转换后的x方向向量（米）
                - left_vec: 转换后的y方向向量（米）
                - z_vec: 转换后的z方向向量（不使用）
                - bottom_angle: bottom向量与x轴的夹角（不使用）
                - left_angle: left向量与y轴的夹角（不使用）
            ax: Matplotlib轴对象
            label_prefix: 标签前缀（标准或测试）
            color: 矩形颜色 (RGB格式, 默认绿色，范围[0,1])
            thickness: 矩形线条粗细
        """
        if vector is None:
            return
        # 从向量中提取中心点坐标和向量（单位为米）
        center_x, center_y = vector['center'][0], vector['center'][1]
        left_vec = vector['left_vec'][:2]
        bottom_vec = vector['bottom_vec'][:2]
        
        # 归一化向量并缩放到固定尺寸（宽0.8米，长2.5米）
        rect_width, rect_height = self.rect_size  # 宽0.8米，长2.5米
        lx, ly = left_vec
        len_left = (lx**2 + ly**2)**0.5
        bx, by = bottom_vec
        len_bottom = (bx**2 + by**2)**0.5
        
        # 归一化并缩放left_vec到2.5米（长度）
        if len_left > 1e-5:
            left_vec = (lx / len_left * rect_height, ly / len_left * rect_height)
        else:
            left_vec = (0, rect_height)
        # 归一化并缩放bottom_vec到0.8米（宽度）
        if len_bottom > 1e-5:
            bottom_vec = (bx / len_bottom * rect_width, by / len_bottom * rect_width)
        else:
            bottom_vec = (rect_width, 0)
        
        # 应用像素缩放（1米 = scale_factor像素）
        center_x, center_y = center_x * self.scale_factor, center_y * self.scale_factor
        left_vec = (left_vec[0] * self.scale_factor, left_vec[1] * self.scale_factor)
        bottom_vec = (bottom_vec[0] * self.scale_factor, bottom_vec[1] * self.scale_factor)
        
        # 计算向量长度（像素）
        lx, ly = left_vec
        len_left = (lx**2 + ly**2)**0.5
        bx, by = bottom_vec
        len_bottom = (bx**2 + by**2)**0.5
        
        center = (center_x, center_y)  # 矩形左下角点（像素）
        # 计算与left_vec垂直的向量
        if len_left > 1e-5:
            perp1 = (ly / len_left, -lx / len_left)  # 顺时针旋转90°
            perp2 = (-ly / len_left, lx / len_left)  # 逆时针旋转90°
            dot1 = bx * perp1[0] + by * perp1[1]
            dot2 = bx * perp2[0] + by * perp2[1]
            if dot1 > dot2:
                adjusted_bottom_vec = (perp1[0] * len_bottom, perp1[1] * len_bottom)
            else:
                adjusted_bottom_vec = (perp2[0] * len_bottom, perp2[1] * len_bottom)
        else:
            adjusted_bottom_vec = bottom_vec

        # 计算矩形的四个角点（像素）
        bottom_left = center
        bottom_right = (center_x + adjusted_bottom_vec[0], center_y + adjusted_bottom_vec[1])
        top_left = (center_x + left_vec[0], center_y + left_vec[1])
        top_right = (top_left[0] + adjusted_bottom_vec[0], top_left[1] + adjusted_bottom_vec[1])
        
        # 调试：打印角点坐标
        print(f"{label_prefix}矩形角点: 左下={bottom_left}, 右下={bottom_right}, 左上={top_left}, 右上={top_right}")
        
        # 绘制矩形
        from matplotlib.patches import Polygon
        pts = np.array([bottom_left, bottom_right, top_right, top_left])
        polygon = Polygon(pts, closed=True, edgecolor=color, linewidth=thickness, fill=False, label=f"{label_prefix}矩形")
        ax.add_patch(polygon)
        
        # 标记四个角点
        ax.scatter(bottom_left[0], bottom_left[1], c='red', s=50, label=f"{label_prefix}左下")
        ax.scatter(bottom_right[0], bottom_right[1], c='green', s=50, label=f"{label_prefix}右下")
        ax.scatter(top_left[0], top_left[1], c='blue', s=50, label=f"{label_prefix}左上")
        ax.scatter(top_right[0], top_right[1], c='yellow', s=50, label=f"{label_prefix}右上")
        
        # 绘制left_vec（蓝色箭头）
        ax.arrow(center_x, center_y, left_vec[0] * 0.5, left_vec[1] * 0.5, 
                 color='blue', width=thickness/2, head_width=thickness*5, head_length=thickness*10, 
                 label=f"{label_prefix}left_vec")
        
        # 绘制adjusted_bottom_vec（红色箭头）
        ax.arrow(center_x, center_y, adjusted_bottom_vec[0] * 0.5, adjusted_bottom_vec[1] * 0.5, 
                 color='red', width=thickness/2, head_width=thickness*5, head_length=thickness*10, 
                 label=f"{label_prefix}bottom_vec")
        
    def create_visualization_image(self, standard_vector, test_vector, min_canvas_size=(500, 500)):
        """
        创建标准和测试位置的矩形可视化图像
        
        Args:
            standard_vector: 标准位置向量 (米)
            test_vector: 测试位置向量 (米)
            min_canvas_size: 最小画布大小 (宽, 高)，默认为(500, 500)像素
            
        Returns:
            OpenCV格式的图像 (BGR)
        """
        # 计算画布大小
        rect_width, rect_height = self.rect_size
        max_vec_length = max(rect_width, rect_height) * self.scale_factor * 2
        canvas_width = max(min_canvas_size[0], int(max_vec_length + 100))
        canvas_height = max(min_canvas_size[1], int(max_vec_length + 150))
        
        # 创建Matplotlib图形
        fig, ax = plt.subplots(figsize=(canvas_width/100, canvas_height/100))
        
        # 绘制标准和测试矩形
        self.draw_rectangle(standard_vector, ax, label_prefix="标准", color=(0, 1, 0), thickness=2)
        self.draw_rectangle(test_vector, ax, label_prefix="测试", color=(0, 0, 1), thickness=2)
        
        # 添加比例尺文字
        scale_length_pixels = 100
        original_scale_length = scale_length_pixels / self.scale_factor
        scale_label = f"比例尺: {scale_length_pixels} 像素 = {original_scale_length:.2f} 米"
        ax.text(0.05, 0.95, scale_label, transform=ax.transAxes, fontsize=12,
                verticalalignment='top', bbox=dict(boxstyle='round', facecolor='white', alpha=0.5))
        
        # 设置坐标轴范围（Y轴上正下负）
        ax.set_xlim(-100, canvas_width)
        ax.set_ylim(-100, canvas_height)
        ax.set_aspect('equal')
        ax.legend()
        
        # 将Matplotlib图形转换为OpenCV图像
        fig.canvas.draw()
        img = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
        img = img.reshape(fig.canvas.get_width_height()[::-1] + (3,))
        plt.close(fig)  # 关闭图形释放内存
        
        # 转换为OpenCV的BGR格式
        img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        
        # 在调试模式下保存中间图像
        if DEBUG_MODE:
            debug_path = os.path.join(self.output_dir, f"debug_{time.strftime('%Y%m%d_%H%M%S')}.png")
            cv2.imwrite(debug_path, img_bgr)
            print(f"调试图像已保存至: {debug_path}")
        
        return img_bgr

def get_screen_resolution():
    """获取连接的HDMI屏幕分辨率"""
    try:
        # 尝试使用xrandr获取分辨率
        output = subprocess.check_output(['xrandr']).decode('utf-8')
        for line in output.splitlines():
            if '*' in line:
                parts = line.split()
                for part in parts:
                    if 'x' in part and '*' in part:
                        resolution = part.split('*')[0]
                        width, height = map(int, resolution.split('x'))
                        return width, height
    except:
        pass
    
    # 默认返回常见分辨率
    return 1920, 1080

def main():
    try:
        # 模型路径
        seg_model_path = "./merge_seg_epoch201_rknn_model"
        pose_model_path = "./merge_pose_epoch600_rknn_model"
        
        print(f"使用分割模型: {seg_model_path}")
        print(f"使用姿态模型: {pose_model_path}")

        # 标准位置图像路径
        standard_image_path = "data/Color1/440.png"
        print(f"标准位置图像路径: {standard_image_path}")
        visualizer = PositionVisualizer(seg_model_path, pose_model_path, standard_image_path, 
                                       rect_size=(0.8, 2.5), scale_factor=100)

        # 获取标准位置向量
        standard_vector = visualizer.calculate_position(standard_image_path)
        if standard_vector is None:
            print("错误: 无法获取标准位置向量")
            return
        print(f"标准位置向量: {standard_vector}")

        # 设置OpenCV全屏窗口
        window_name = "Position Visualization"
        cv2.namedWindow(window_name, cv2.WND_PROP_FULLSCREEN)
        cv2.setWindowProperty(window_name, cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)
        
        # 获取屏幕分辨率
        screen_width, screen_height = get_screen_resolution()
        print(f"检测到屏幕分辨率: {screen_width}x{screen_height}")

        # 测试图像范围：200.png 到 440.png
        start_idx = 1
        end_idx = 500
        test_image_dir = "data/Color1"
        
        for idx in range(start_idx, end_idx + 1):
            test_image_path = os.path.join(test_image_dir, f"{idx}.png")
            if not os.path.exists(test_image_path):
                print(f"测试图像 {test_image_path} 不存在，跳过")
                continue
            
            print(f"\n处理测试图像: {test_image_path}")
            
            # 获取测试位置向量
            test_vector = visualizer.calculate_position(test_image_path)
            if test_vector is None:
                print(f"警告: 无法获取测试位置向量 {test_image_path}")
                continue
            print(f"测试位置向量: {test_vector}")
            
            # 创建可视化图像
            img_bgr = visualizer.create_visualization_image(standard_vector, test_vector)
            
            # 在调试模式下显示图像尺寸信息
            if DEBUG_MODE:
                print(f"可视化图像尺寸: {img_bgr.shape[1]}x{img_bgr.shape[0]}")
            
            # 调整图像大小以适应屏幕
            scale_factor = min(screen_width / img_bgr.shape[1], screen_height / img_bgr.shape[0])
            if scale_factor < 1:
                img_bgr = cv2.resize(img_bgr, (0, 0), fx=scale_factor, fy=scale_factor)
            
            # 居中显示
            pad_x = (screen_width - img_bgr.shape[1]) // 2
            pad_y = (screen_height - img_bgr.shape[0]) // 2
            padded_img = np.zeros((screen_height, screen_width, 3), dtype=np.uint8)
            padded_img[pad_y:pad_y+img_bgr.shape[0], pad_x:pad_x+img_bgr.shape[1]] = img_bgr
            
            # 添加标题
            title = f"标准位置 (440.png) vs 测试位置 ({idx}.png)"
            cv2.putText(padded_img, title, (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
            
            # 显示图像
            cv2.imshow(window_name, padded_img)
            
            # 保存结果
            output_path = os.path.join(visualizer.output_dir, f"position_comparison_{idx}.png")
            cv2.imwrite(output_path, padded_img)
            print(f"结果已保存至: {output_path}")
            
            # 等待按键或短暂延迟
            key = cv2.waitKey(100)  # 100毫秒延迟
            if key == 27:  # ESC键退出
                break
        
        print("所有测试图像处理完成")
        
    except Exception as e:
        import traceback
        print(f"执行过程中发生错误: {e}")
        print(traceback.format_exc())
        
        # 在调试模式下保存错误信息
        if DEBUG_MODE:
            error_path = os.path.join(visualizer.output_dir, f"error_{time.strftime('%Y%m%d_%H%M%S')}.txt")
            with open(error_path, 'w') as f:
                f.write(f"错误: {e}\n")
                f.write(traceback.format_exc())
            print(f"错误信息已保存至: {error_path}")
        
    finally:
        # 确保关闭窗口
        cv2.destroyAllWindows()

if __name__ == "__main__":
    # 添加命令行参数支持
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--debug', action='store_true', help='启用调试模式')
    args = parser.parse_args()
    
    if args.debug:
        os.environ['DEBUG_MODE'] = '1'
    
    main()